package rabbit.flt.plugins.common.plugin;

import rabbit.flt.common.context.TraceContext;
import rabbit.flt.common.context.TraceMetaData;
import rabbit.flt.common.trace.MethodStackInfo;
import rabbit.flt.common.trace.TraceData;
import rabbit.flt.plugins.common.AbstractPlugin;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.lang.reflect.Method;

import static rabbit.flt.common.trace.MessageType.METHOD;

/**
 * 性能插件，统计开销
 */
public class ReactorPerformancePlugin extends AbstractPlugin {

    @Override
    protected void fillTraceData(TraceData traceData, Object objectEnhanced, Method method, Object[] args, Object result) {
        traceData.setMessageType(METHOD.name());
    }

    @Override
    public Object[] before(Object objectEnhanced, Method method, Object[] args) {
        if (isTraceOpened()) {
            return super.before(objectEnhanced, method, args);
        }
        return args;
    }

    /**
     * 判断是不是reactor方法
     *
     * @param method
     * @return
     */
    protected final boolean isReactor(Method method) {
        Class<?> returnType = method.getReturnType();
        return Mono.class.isAssignableFrom(returnType) || Flux.class.isAssignableFrom(returnType);
    }

    @Override
    public Object after(Object objectEnhanced, Method method, Object[] args, Object result) {
        if (isTraceOpened() && isReactor(method)) {
            TraceMetaData traceMetaData = TraceContext.getTraceMetaData();
            MethodStackInfo stackInfo = getStackInfo(method);
            TraceData traceData = stackInfo.getTraceData();
            if (result instanceof Mono) {
                return ((Mono<?>) result).onErrorResume(e -> {
                    saveError2Trace(method, traceMetaData, traceData, e);
                    return Mono.error(e);
                }).doOnSuccess(v -> traceData.updateCost()).map(v -> {
                    traceData.updateCost();
                    return v;
                }).doOnCancel(() -> handleCancel(traceData)).doFinally(s -> handleReactorFinal(objectEnhanced, method, args, result, stackInfo, traceData));
            }
            if (result instanceof Flux) {
                return ((Flux<?>) result).doOnCancel(() -> handleCancel(traceData)).onErrorResume(e -> {
                    saveError2Trace(method, traceMetaData, traceData, e);
                    return Mono.error(e);
                }).doOnComplete(traceData::updateCost).map(v -> {
                    traceData.updateCost();
                    return v;
                }).doFinally(s -> handleReactorFinal(objectEnhanced, method, args, result, stackInfo, traceData));
            }
        }
        return result;
    }

    /**
     * 处理cancel事件
     *
     * @param traceData
     */
    protected void handleCancel(TraceData traceData) {
        traceData.updateCost();
    }

    /**
     * 处理reactor final调用
     *
     * @param objectEnhanced
     * @param method
     * @param args
     * @param result
     * @param stackInfo
     * @param traceData
     */
    protected void handleReactorFinal(Object objectEnhanced, Method method, Object[] args, Object result, MethodStackInfo stackInfo, TraceData traceData) {
        if (stackInfo.isPopped()) {
            fillTraceData(traceData, objectEnhanced, method, args, result);
            super.handleTraceData(traceData);
        }
    }

    /**
     * 将异常信息保存到trace信息
     *
     * @param method
     * @param traceMetaData
     * @param traceData
     * @param e
     */
    protected void saveError2Trace(Method method, TraceMetaData traceMetaData, TraceData traceData, Throwable e) {
        traceMetaData.trySaveErrorMethod(method);
        traceData.setStatus(TraceData.Status.ERR);
        // 更新开销
        traceData.updateCost();
        if (traceMetaData.isErrorMethod(method)) {
            traceData.setData(getStackText(e));
        }
    }

    /**
     * 返回值如果是Mono或者Flux则不发送数据
     *
     * @param traceData
     */
    @Override
    protected void handleTraceData(TraceData traceData) {
        // reactor 类型的发送在 final环节处理
        if (!isReactor(traceData.getTargetMethod())) {
            super.handleTraceData(traceData);
        }
    }

    @Override
    public void doFinal(Object objectEnhanced, Method method, Object[] args, Object result) {
        if (isTraceOpened()) {
            super.doFinal(objectEnhanced, method, args, result);
        }
    }

    @Override
    public void onException(Object objectEnhanced, Method method, Object[] args, Throwable t) {
        // reactor 类型的异常在 sendTrace中发送
        if (isTraceOpened() && !isReactor(method)) {
            super.onException(objectEnhanced, method, args, t);
        }
    }

}
